#pragma once

#include <stdint.h>

#include <unordered_map>
#include <unordered_set>
#include <vector>

#include "hawking/util/node_content.pb.h"
#include "util/arena/arena_wrapper.h"
#include "doc/normal_document.h"

namespace hawking {
namespace indexlib {

struct DocContext {
    uint64_t pk;
    int32_t filtered = 0;
    // string and array
    std::unordered_map<int32_t, util::FieldValue*> field2attr;
    // single value
    std::unordered_map<int32_t, int64_t> single_int_field;
    std::unordered_map<int32_t, float> single_float_field;
    std::unordered_map<int32_t, double> single_double_field;

    util::FieldValue* GetField(int32_t field_id) const {
        auto iter = field2attr.find(field_id);
        return iter != field2attr.end() ? iter->second : nullptr;
    }

    bool GetField(int32_t field_id, int64_t* ret) const {
        auto iter = single_int_field.find(field_id);
        if (iter != single_int_field.end()) {
            *ret = iter->second;
            return true;
        }

        return false;
    }

    bool GetField(int32_t field_id, float* ret) const {
        auto iter = single_float_field.find(field_id);
        if (iter != single_float_field.end()) {
            *ret = iter->second;
            return true;
        }

        return false;
    }

    bool GetField(int32_t field_id, double* ret) const {
        auto iter = single_double_field.find(field_id);
        if (iter != single_double_field.end()) {
            *ret = iter->second;
            return true;
        }

        return false;
    }

    DocContext() = default;
};

class SegmentContext {
public:
    SegmentContext() = default;
    SegmentContext(indexlib::ArenaWrapper& aw) : arena_wrapper_(&aw) {}
    ~SegmentContext() = default;

    uint32_t RelateDocCount() const {
        return relate_;
    }

    void DecrRelateCount() {
        --relate_;
    }

    void SetRelateDocCount(uint32_t r) {
        relate_ = r;
    }

    std::vector<DocId>* MutableDocIds() {
        return &doc_ids_;
    }    

    const std::vector<DocId>& DocIds() const {
        return doc_ids_;
    }

    DocId* MutableDocIds() {
        return (DocId*)doc_ids_.data();
    }

    void PushDocId(DocId doc_id) {
        doc_ids_.emplace_back(doc_id);
        docs_ctxs_.emplace_back();
        ++relate_;
    }

    const std::vector<DocContext>& DocContexts() const {
        return docs_ctxs_;
    }

    std::vector<DocContext>* MutableDocContext() {
        return &docs_ctxs_;
    }

    void Clear() {
        docs_ctxs_.clear();
        doc_ids_.clear();
        relate_ = 0;
    }

    void Reserve(size_t count) {
        docs_ctxs_.reserve(count);
        doc_ids_.reserve(count);
    }

    indexlib::ArenaWrapper& ArenaCtx() {
        return *arena_wrapper_;
    }

    void InsertVisitedFieldId(int32_t field_id) {
        visited_field_ids_.insert(field_id);
    }

    bool IsVisitedFieldId(int32_t field_id) const {
        return visited_field_ids_.find(field_id) != visited_field_ids_.end();
    }

private:
    uint32_t relate_ = 0;
    indexlib::ArenaWrapper* arena_wrapper_ = nullptr;

    // separate the docs and ctxs for avoid copy docs one more time when read inverted
    std::vector<DocContext> docs_ctxs_;
    std::vector<DocId> doc_ids_;

    std::unordered_set<int32_t> visited_field_ids_;
};

}
}